"""
Base class for vision backbones in continual learning.

This module defines the interface that all backbone networks must implement to be used
within the continual learning framework. Backbones are responsible for feature extraction
and can be adapted for continual learning using various strategies like prompting,
adapters, or partial fine-tuning.

A backbone can be any neural network architecture (CNN, ViT, etc.) that processes
input data and produces feature representations suitable for downstream tasks.
"""

from abc import abstractmethod
import torch

from continuallearning.interfaces import BackboneInterface
from continuallearning.interfaces import BackboneOutput
from ..core.output_utils import standardize_backbone_output


class BaseBackbone(BackboneInterface):
    """
    Base class for all backbone networks in continual learning.

    Backbones are responsible for feature extraction from input data and serve as the
    foundation for continual learning models. They can be pre-trained networks that are
    adapted for continual learning through various approaches:

    - Fixed feature extractors with task-specific heads
    - Parameter-efficient tuning (prompts, adapters)
    - Selective fine-tuning of specific layers
    - Full fine-tuning with regularization

    All backbone models should inherit from this class and implement the forward method
    that returns a standardized output, making it flexible for different architectures
    and downstream components.
    """

    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor, **kwargs) -> BackboneOutput:
        """
        Extract features from input data.

        Subclasses must implement this method to perform the actual feature extraction.
        The implementation should return either a BackboneOutput instance or a dict
        that can be converted to a BackboneOutput.

        Args:
            x: Input tensor
            **kwargs: Additional arguments for flexibility

        Returns:
            Union[Dict[str, Any], BackboneOutput]: Output containing features and optional
                                           hidden states or attention weights
        """
        # This is an abstract method that must be implemented by subclasses
        raise NotImplementedError("Subclasses must implement forward method")

    @abstractmethod
    def get_feature_dims(self) -> int:
        """
        Get the feature dimensions output by the backbone.

        This method returns the dimensionality of the feature representations
        produced by the backbone, which is essential for constructing compatible
        classification heads or adapter modules.

        Returns:
            int: Feature dimensions (e.g., 768 for ViT-B, 2048 for ResNet-50)
        """
        raise NotImplementedError

    def extract_features(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
        """
        Extract features from input without additional processing.

        This implementation ensures consistent behavior across backbone models
        by using the standardization utility.

        Args:
            x: Input tensor
            **kwargs: Additional arguments for the backbone

        Returns:
            torch.Tensor: Feature tensor representation
        """
        output = self.forward(x, **kwargs)
        standardized = standardize_backbone_output(output)

        if standardized.features is None:
            raise ValueError("Backbone did not produce valid features")

        return standardized.features
